# -*- coding: utf-8 -*-
# 会话生成器基类，定义公共接口

import os
import pandas as pd
from abc import ABC, abstractmethod
from core.data_loader import load_excel_data, load_csv_data

class SessionGenerator(ABC):
    """会话生成器基类"""
    
    def __init__(self, template_path=None):
        self.template_path = template_path
        self.template_content = None
        if template_path and os.path.exists(template_path):
            self.load_template(template_path)
    
    def load_template(self, template_path, encoding='utf-8'):
        """加载模板文件"""
        try:
            with open(template_path, 'r', encoding=encoding) as f:
                self.template_content = f.read()
            return True
        except Exception as e:
            print(f"加载模板文件失败: {e}")
            return False
    
    def load_data(self, data_path, sheet_name='Sheet1'):
        """加载数据文件 (Excel/CSV)"""
        try:
            if data_path.endswith('.csv'):
                data, error = load_csv_data(data_path)
            else:  # 默认为Excel
                data, error = load_excel_data(data_path, sheet_name)
            
            if error:
                print(f"加载数据文件失败: {error}")
                return [], error
            
            if data is None or len(data) == 0:
                return [], "数据文件中没有可用记录"
            
            # 确保返回的是字典列表
            records = data.to_dict(orient='records')
            
            # 验证每条记录的格式
            for i, record in enumerate(records):
                if not isinstance(record, dict):
                    print(f"警告: 第{i+1}条记录不是字典格式，将被跳过")
                    continue
                    
                # 确保必要字段存在
                if 'hostname' not in record or 'IP' not in record or 'role' not in record:
                    print(f"警告: 第{i+1}条记录缺少必要字段，将被跳过")
                    continue
            
            # 打印加载的数据进行调试
            print(f"成功加载 {len(records)} 条记录:")
            for i, record in enumerate(records):
                print(f"  记录 {i+1}: hostname={record.get('hostname', '无')}, "
                      f"IP={record.get('IP', '无')}, role={record.get('role', '无')}")
            
            return records, None
        except Exception as e:
            import traceback
            error_msg = f"加载数据文件异常: {str(e)}"
            print(error_msg)
            print(f"详细错误: {traceback.format_exc()}")
            return [], error_msg
    
    @abstractmethod
    def generate_session(self, machine_data, output_dir='.'):
        """生成会话文件 (需要子类实现)"""
        pass
    
    def generate_all_sessions(self, data_path, output_dir=None, sheet_name='Sheet1'):
        """生成所有会话文件"""
        if not self.template_content:
            print("错误: 模板文件未加载")
            return False
        
        # 加载数据 - 确保检查错误
        machines, error = self.load_data(data_path, sheet_name)
        if error:
            print(f"加载数据文件失败: {error}")
            return False
        
        if not machines or len(machines) == 0:
            print("错误: 机器列表为空")
            return False
        
        # 确保输出目录存在
        client_type = self.__class__.__name__.lower().replace('generator', '')
        base_dir = self.get_output_path(client_type, output_dir)
        
        # 为每台机器生成会话文件
        success_count = 0
        for idx, machine in enumerate(machines):
            print(f"正在处理第 {idx+1}/{len(machines)} 台机器...")
            if isinstance(machine, dict) and self.generate_session(machine, output_dir):
                success_count += 1
            else:
                print(f"跳过第 {idx+1} 台机器")
        
        print(f"已成功生成 {success_count}/{len(machines)} 个会话文件，保存在 {base_dir}")
        return success_count > 0

    def get_output_path(self, client_type, base_output_dir=None, role=None):
        """获取统一的输出路径
        
        Args:
            client_type: 客户端类型 (mobaxterm, securecrt, xshell)
            base_output_dir: 用户指定的基础输出目录，None则使用默认值
            role: 服务器角色/分类
            
        Returns:
            str: 完整的输出路径
        """
        # 确定基础目录
        if base_output_dir is None or not base_output_dir:
            # 如果未指定，则使用当前目录
            base_dir = os.getcwd()
        else:
            # 使用用户指定的目录
            base_dir = base_output_dir
        
        # 在基础目录下创建sessions文件夹
        sessions_dir = os.path.join(base_dir, 'sessions')
        os.makedirs(sessions_dir, exist_ok=True)
        
        # 按客户端类型创建子目录
        client_dir = os.path.join(sessions_dir, client_type)
        os.makedirs(client_dir, exist_ok=True)
        
        # 如果指定了角色，则创建角色子目录
        if role:
            role_dir = os.path.join(client_dir, role)
            os.makedirs(role_dir, exist_ok=True)
            return role_dir
        
        return client_dir

    def check_directory_writeable(self, directory):
        """检查目录是否可写"""
        if not os.path.exists(directory):
            try:
                os.makedirs(directory, exist_ok=True)
            except Exception as e:
                print(f"创建目录失败: {directory}, 错误: {e}")
                return False
        
        # 测试是否可写
        test_file = os.path.join(directory, '.write_test')
        try:
            with open(test_file, 'w') as f:
                f.write('test')
            os.remove(test_file)
            return True
        except Exception as e:
            print(f"目录不可写: {directory}, 错误: {e}")
            return False 